import random
import numpy as np
import pandas as pd
from deap import base, creator, tools, algorithms

# 读取用户上传的数据
file_path = '附件1.xlsx'
df_land = pd.read_excel(file_path, sheet_name='乡村的现有耕地')

file_path_2 = '附件2.xlsx'
df_crops = pd.read_excel(file_path_2, sheet_name='2023年统计的相关数据')

# 定义豆类作物的编号集合
bean_crops = {1, 2, 3, 4, 5, 17, 18, 19}

# 整理地块类型并分类
A = {row['地块名称']: row['地块面积/亩'] for _, row in df_land.iterrows()}

# 单季地块：平旱地、梯田、山坡地
plots_single_season = df_land[df_land['地块类型'].isin(['平旱地', '梯田', '山坡地'])]['地块名称'].tolist()

# 双季地块：水浇地、普通大棚、智慧大棚
plots_double_season = df_land[df_land['地块类型'].isin(['水浇地', '普通大棚', '智慧大棚'])]['地块名称'].tolist()

# 将双季节再进行细分一下
irrigated_land = df_land[df_land['地块类型'] == '水浇地']['地块名称'].tolist()
greenhouse_land = df_land[df_land['地块类型'] == '普通大棚']['地块名称'].tolist()
smart_greenhouse_land = df_land[df_land['地块类型'] == '智慧大棚']['地块名称'].tolist()

# 定义作物类型
crops_single_season = list(range(1, 16))  # 作物1-15，适用于单季地块
crops_double_season_irrigated = list(range(16, 38))  # 作物16-37，适用于水浇地
crops_double_season_greenhouse = list(range(23, 42))  # 作物23-41，适用于普通大棚
crops_double_season_smart = list(range(27, 35))  # 作物27-34，适用于智慧大棚

# 处理销售单价列，取范围的平均值
df_crops['销售单价(元/斤)'] = df_crops['销售单价/(元/斤)'].apply(
    lambda x: (float(x.split('-')[0]) + float(x.split('-')[1])) / 2 if isinstance(x, str) else x
)

# 生成销售单价、亩产量和种植成本的字典
P = {row['作物编号']: row['销售单价(元/斤)'] for _, row in df_crops.iterrows()}  # 销售单价
Y = {row['作物编号']: row['亩产量/斤'] for _, row in df_crops.iterrows()}  # 亩产量
C = {row['作物编号']: row['种植成本/(元/亩)'] for _, row in df_crops.iterrows()}  # 种植成本

# 定义年份范围
years = list(range(2024, 2031))  # 从2024到2030
num_years = len(years)

# 获取所有地块
plots_all = plots_single_season + irrigated_land + greenhouse_land + smart_greenhouse_land

# DEAP 库配置
creator.create("FitnessMax", base.Fitness, weights=(1.0,))
creator.create("Individual", list, fitness=creator.FitnessMax)

toolbox = base.Toolbox()

# 基因初始化：每个基因可以是 0, 0.5 或 1
def random_gene():
    return random.choice([0, 0.5, 1])

# 初始化个体
toolbox.register("attr_float", random_gene)
toolbox.register("individual", tools.initRepeat, creator.Individual, toolbox.attr_float, 
                 n=len(plots_all) * len(crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart) * num_years * 2)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

# 适应度函数
def evaluate(individual):
    solution = np.reshape(individual, (len(plots_all), len(crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart), num_years, 2))
    total_profit = 0
    
    # 计算收益
    for i, plot in enumerate(plots_all):
        for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
            for t_idx, t in enumerate(years):
                for j in [0, 1]:
                    planting_ratio = solution[i][k-1][t_idx][j]
                    total_profit += planting_ratio * (P[k] * Y[k] * A[plot] - C[k] * A[plot])
    
    penalty = 0

    # 约束1：种植比例之和为1
    for i, plot in enumerate(plots_all):
        for t_idx, t in enumerate(years):
            for j in [0, 1]:
                total_ratio = sum(solution[i][k-1][t_idx][j] for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart)
                if abs(total_ratio - 1) > 0.01:
                    penalty += 1000  # 如果总种植比例不为1，则扣分

    # 约束2：每三年种植一次豆类作物
    for i, plot in enumerate(plots_all):
        for t_start in range(num_years - 3 + 1):
            bean_planted = any(solution[i][k-1][t_idx][j] > 0 for k in bean_crops for t_idx in range(t_start, t_start + 3) for j in [0, 1])
            if not bean_planted:
                penalty += 1000  # 如果未满足轮作约束则扣分
    
    return total_profit - penalty,

# 注册适应度函数、选择、交叉和变异
toolbox.register("evaluate", evaluate)
toolbox.register("mate", tools.cxTwoPoint)
toolbox.register("mutate", tools.mutFlipBit, indpb=0.05)
toolbox.register("select", tools.selTournament, tournsize=3)

# 遗传算法主程序
def main():
    pop = toolbox.population(n=50)  # 种群大小
    hof = tools.HallOfFame(1)  # 保存最佳个体
    stats = tools.Statistics(lambda ind: ind.fitness.values)
    stats.register("avg", np.mean)
    stats.register("min", np.min)
    stats.register("max", np.max)

    algorithms.eaSimple(pop, toolbox, cxpb=0.7, mutpb=0.2, ngen=100, stats=stats, halloffame=hof, verbose=True)

    return hof[0]  # 返回最佳个体

# 运行遗传算法
best_individual = main()

# 解码最佳解并保存为 Excel
solution = np.reshape(best_individual, (len(plots_all), len(crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart), num_years, 2))
results = []

for i, plot in enumerate(plots_all):
    for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
        for t_idx, t in enumerate(years):
            for j in [0, 1]:
                planting_ratio = solution[i][k-1][t_idx][j]
                if planting_ratio > 0:
                    results.append([plot, k, t, j+1, planting_ratio])

# 将结果保存为 DataFrame 并写入 Excel 文件
df_results = pd.DataFrame(results, columns=["地块", "作物编号", "年份", "季节", "种植比例"])
output_path = "optimization_results_deap.xlsx"
df_results.to_excel(output_path, index=False)
print("结果已保存为 Excel 文件:", output_path)